#!/usr/bin/env python3
"""
主程序：PBPK（生理药代动力学）建模系统

该程序整合了吸收、分布和消除模块，使用人群数据作为输入，
完成整个PBPK建模过程并生成可视化图表。

作者：BlackCat@CPPO
版本：1.0.0
日期：2025-07-22
"""

import os
import sys
import argparse
import json
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime

# 添加模块路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from src.absorption.absorption_module import AbsorptionModel
from src.distribution.distribution_module import DistributionModel
from src.elimination.elimination_module import EliminationModel
from src.population.population_data import PopulationData
from src.visualization.plot_generator import PlotGenerator
from src.utils.config_manager import ConfigManager
from src.utils.logger import PBPKLogger as Logger

class PBPKModelingSystem:
    """
    主PBPK建模系统类
    
    该类协调所有PBPK建模组件，包括吸收、分布、消除和人群数据处理。
    """
    
    def __init__(self, config_path=None):
        """
        初始化PBPK建模系统
        
        参数:
            config_path (str): 配置文件路径，默认为None使用默认配置
        """
        self.config = ConfigManager(config_path)
        self.logger = Logger(config=self.config.config.get('logging', {}))
        
        # 初始化各模块
        self.absorption = None
        self.distribution = None
        self.elimination = None
        self.population = None
        self.visualizer = None
        
        self.logger.info("PBPK建模系统初始化完成")
    
    def load_population_data(self, population_file=None):
        """
        加载人群数据
        
        参数:
            population_file (str): 人群数据文件路径
            
        返回:
            dict: 人群数据字典
        """
        self.population = PopulationData()
        
        if population_file:
            data = self.population.load_from_file(population_file)
        else:
            # 使用默认人群数据
            data = self.population.get_default_population()
        
        self.logger.info(f"人群数据加载完成，包含{len(data)}个个体")
        return data
    
    def setup_models(self, drug_params=None):
        """
        设置PBPK模型参数
        
        参数:
            drug_params (dict): 药物特定参数
        """
        # 获取基础配置
        base_params = self.config.get_model_params()
        
        # 合并药物参数
        if drug_params:
            base_params.update(drug_params)
        
        # 初始化各模块
        self.absorption = AbsorptionModel(base_params)
        self.distribution = DistributionModel(base_params)
        self.elimination = EliminationModel(base_params)
        self.visualizer = PlotGenerator()
        
        self.logger.info("PBPK模型参数设置完成")
    
    def run_simulation(self, population_data, simulation_time=24, time_step=0.1):
        """
        运行PBPK模拟
        
        参数:
            population_data (dict): 人群数据
            simulation_time (float): 模拟时间（小时）
            time_step (float): 时间步长（小时）
            
        返回:
            dict: 模拟结果
        """
        self.logger.info(f"开始PBPK模拟，模拟时间：{simulation_time}小时")
        
        # 生成时间序列
        time_points = np.arange(0, simulation_time + time_step, time_step)
        
        # 存储所有个体的结果
        all_results = {}
        
        for individual_id, individual_data in population_data.items():
            self.logger.info(f"处理个体：{individual_id}")
            
            # 更新个体特定参数
            individual_params = self.config.get_individual_params(individual_data)
            
            # 运行模拟
            result = self._simulate_individual(
                individual_id, 
                individual_params, 
                time_points
            )
            
            all_results[individual_id] = result
        
        self.logger.info("PBPK模拟完成")
        return all_results
    
    def _simulate_individual(self, individual_id, params, time_points):
        """
        模拟单个个体
        
        参数:
            individual_id (str): 个体ID
            params (dict): 个体参数
            time_points (np.array): 时间点数组
            
        返回:
            dict: 个体模拟结果
        """
        # 计算吸收阶段
        absorption_result = self.absorption.calculate_absorption(
            params, time_points
        )
        
        # 计算分布阶段
        distribution_result = self.distribution.calculate_distribution(
            params, time_points, absorption_result
        )
        
        # 计算消除阶段
        elimination_result = self.elimination.calculate_elimination(
            params, time_points, distribution_result
        )
        
        # 整合结果
        result = {
            'time': time_points,
            'absorption': absorption_result,
            'distribution': distribution_result,
            'elimination': elimination_result,
            'concentration': elimination_result['plasma_concentration'],
            'parameters': params
        }
        
        return result
    
    def generate_plots(self, results, output_dir='../output'):
        """
        生成可视化图表
        
        参数:
            results (dict): 模拟结果
            output_dir (str): 输出目录
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 生成各类图表
        plots = []
        
        # 1. 浓度-时间曲线
        concentration_plot = self.visualizer.plot_concentration_time(
            results, 
            os.path.join(output_dir, f'concentration_time_{timestamp}.png')
        )
        plots.append(concentration_plot)
        
        # 2. 吸收分布消除过程
        process_plot = self.visualizer.plot_adme_processes(
            results,
            os.path.join(output_dir, f'adme_processes_{timestamp}.png')
        )
        plots.append(process_plot)
        
        # 3. 人群变异性分析
        variability_plot = self.visualizer.plot_population_variability(
            results,
            os.path.join(output_dir, f'population_variability_{timestamp}.png')
        )
        plots.append(variability_plot)
        
        # 4. 参数敏感性分析
        sensitivity_plot = self.visualizer.plot_parameter_sensitivity(
            results,
            os.path.join(output_dir, f'sensitivity_analysis_{timestamp}.png')
        )
        plots.append(sensitivity_plot)
        
        self.logger.info(f"图表生成完成，保存在：{output_dir}")
        return plots
    
    def export_results(self, results, output_dir='../output'):
        """
        导出模拟结果
        
        参数:
            results (dict): 模拟结果
            output_dir (str): 输出目录
        """
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        
        # 导出为CSV
        for individual_id, result in results.items():
            df = pd.DataFrame({
                'time': result['time'],
                'plasma_concentration': result['concentration'],
                'absorption_rate': result['absorption']['rate'],
                'distribution_volume': result['distribution']['volume'],
                'clearance_rate': result['elimination']['clearance']
            })
            
            filename = os.path.join(output_dir, f'results_{individual_id}_{timestamp}.csv')
            df.to_csv(filename, index=False)
        
        # 导出为JSON
        summary_file = os.path.join(output_dir, f'simulation_summary_{timestamp}.json')
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, indent=2, default=str)
        
        self.logger.info(f"结果导出完成，保存在：{output_dir}")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='PBPK建模系统')
    parser.add_argument('--config', help='配置文件路径')
    parser.add_argument('--population', help='人群数据文件路径')
    parser.add_argument('--drug', help='药物参数文件路径')
    parser.add_argument('--time', type=float, default=24, help='模拟时间（小时）')
    parser.add_argument('--step', type=float, default=0.1, help='时间步长（小时）')
    parser.add_argument('--output', default='../output', help='输出目录')
    
    args = parser.parse_args()
    
    # 初始化系统
    system = PBPKModelingSystem(args.config)
    
    try:
        # 加载人群数据
        population_data = system.load_population_data(args.population)
        
        # 加载药物参数
        drug_params = {}
        if args.drug:
            with open(args.drug, 'r', encoding='utf-8') as f:
                drug_params = json.load(f)
        
        # 设置模型
        system.setup_models(drug_params)
        
        # 运行模拟
        results = system.run_simulation(population_data, args.time, args.step)
        
        # 生成图表
        plots = system.generate_plots(results, args.output)
        
        # 导出结果
        system.export_results(results, args.output)
        
        print("PBPK建模完成！")
        print(f"结果保存在：{args.output}")
        
    except Exception as e:
        print(f"错误：{str(e)}")
        system.logger.error(f"模拟失败：{str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()